
import asyncio
import random
import time
from asyncio import current_task, Queue, create_task
import httpx
from httpx import AsyncClient, Response
from httpx._types import RequestFiles
from numpy import mean
from sqlalchemy import select, text
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine

headers = {
    'User-Agent': 'Mozilla/5.0 (iPhone; CPU iPhone OS 13_2_3 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/13.0.3 Mobile/15E148 Safari/604.1',
}

exec_time_list: list[float] = []

async def get_db_session(database_info: str) -> sessionmaker:
    return sessionmaker(
        class_=AsyncSession,
        autocommit=False,
        autoflush=False,
        bind=create_async_engine(database_info)
    )

def field(func):
    async def wrapper(*params, **kwparams):
        start_time = time.time() * 1000

        result = await func(*params, **kwparams)

        end_time = time.time() * 1000

        exec_time_list.append(end_time - start_time)

        print(result.status_code, " | ", result.json(), " | ", str(int(end_time - start_time) + 1) + "ms")

    return wrapper

@field
async def request(client: AsyncClient, method: str, url: str, param_type: str, params: dict, headers: dict,
                  files: RequestFiles = None) -> Response:
    if param_type == "json":
        result = await client.request(method=method, url=url, json=params, files=files, headers=headers)
    elif param_type == "from":
        result = await client.request(method=method, url=url, data=params, files=files, headers=headers)
    else:
        result = httpx.request(method=method, url=url, headers=headers)

    return result

async def task_handle(client: AsyncClient, queue: Queue):
    token = await queue.get()
    headers['authorization'] = token

    await request(client, method="", url="", param_type="json", params={}, headers=headers)

    await request(client, method="", url="", param_type="from", params={}, headers=headers)

    await request(client, method="", url="", param_type="json", params={}, files={"file": ""}, headers=headers)

    ...

async def task_handle_pool(client: AsyncClient, queue: Queue):
    while True:
        if queue.empty():
            break

        token = await queue.get()
        headers['authorization'] = token

        await request(client, method="", url="", param_type="json", params={}, headers=headers)

        await request(client, method="", url="", param_type="from", params={}, headers=headers)

        await request(client, method="", url="", param_type="json", params={}, files={"file": ""}, headers=headers)

        ...

async def pressure_measurement(database_info: str, sql: str, is_pool: bool):
    queue = Queue()

    async_session = await get_db_session(database_info)
    async with async_session() as work_session:
        users = (await work_session.execute(text(sql))).fetchall()
        start_time = time.time() * 1000
        async with httpx.AsyncClient(timeout=None) as client:
            token_task = []
            for user in users:
                token_task.append(
                    asyncio.create_task(
                        client.post(url="", json={
                            "mobile": user[1],
                            "password": ""
                        }, headers=headers)
                    )
                )

                # result = await client.post(url="", json={
                #     "mobile": user[0],
                #     "password": user[1]
                # }, headers=headers)
                # token_info = result.json()["data"]
                # await queue.put(token_info["token_type"] + " " + token_info["token"])

            await asyncio.wait(token_task)
            end_time = time.time() * 1000
            for t in token_task:
                token_info = t.result().json()["data"]
                await queue.put(token_info["token_type"] + " " + token_info["token"])
            print("queue_len：", queue.qsize(), " | ", str(int(end_time - start_time) + 1) + "ms")

    async with httpx.AsyncClient(timeout=None) as client:
        task_list = []
        if is_pool:
            for t in range(int(queue.qsize()/2)):
                task_list.append(create_task(task_handle_pool(client, queue)))
        else:
            for t in range(queue.qsize()):
                task_list.append(create_task(task_handle(client, queue)))

        await asyncio.wait(task_list)

        average_time = mean(exec_time_list)

        print("最大响应时间：", max(exec_time_list))
        print("最小响应时间：", average_time)
        print("平均应时间：", min(exec_time_list))
        print("QPS：", queue.qsize() / average_time)

if __name__ == '__main__':
    asyncio.get_event_loop().run_until_complete(
        pressure_measurement("", "", False))




























































































